package edu.purdue.cs.ds;

import java.io.*;
import java.math.BigInteger;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Iterator;

import edu.purdue.cs.ds.vss.*;

/**
 * A Server is responsible for a single share
 * of the secret distributed by the dealer.
 * @author jwallrab
 *
 */
public class Server {

	private NetworkCommunicator net;
	private ArrayList<VerifiableShare> myShares;
	private String filename = "files/hosts";
	private String local_hostname;
	private int my_id;
	private int n;
	private int k;
	private int t;
	private int v = 0;
	private final int timeout = 2000;
	private Cryptography crypto;
	private ArrayList<ConnectionHandler> connections;
	private IndexSets iset;
	private ShareStore old_store;
	private ShareStore new_store;
	private ArrayList<String> verifiedServers;
	private int numRefreshes = 1;
	private boolean failures = false; //True if servers should fail
	private boolean nonresponsive = false; //True if server should never respond

	/**
	 * Constructor
	 * @param p Port
	 * @param npar Number of Servers
	 * @param k Number of Shares needed to Reconstruct
	 */
	public Server(int p, int npar, int kpar, boolean evil, boolean nonresponse, int refresh){

		numRefreshes = refresh;
		nonresponsive = nonresponse;
		failures = evil;
		n = npar;
		k = kpar;
		t = k;
		net = new NetworkCommunicator(p);
		connections = new ArrayList<ConnectionHandler>();
		verifiedServers = new ArrayList<String>();
		iset = new IndexSets();
		old_store = new ShareStore();
		new_store = new ShareStore();

		getLocalHostname();
		verifiedServers.add(local_hostname);
		crypto = new Cryptography();
		receiveShare();

		ConnectToAll();
		for(int i = 0; i < numRefreshes; ++i)
			VerifyStage();
	}

	/**
	 * Returns the machine's host name
	 */
	public void getLocalHostname(){
		try {
			local_hostname = java.net.InetAddress.getLocalHost().getHostName();
		} catch (UnknownHostException e) {
			e.printStackTrace();
		}
	}

	/**
	 * Receive Share from Dealer
	 */
	public void receiveShare(){

		try {
			net.openServerSocket();
			net.openConnection();
			int num_shares = net.receiveInt();

			myShares = new ArrayList<VerifiableShare>();

			for(int i = 0; i < num_shares; ++i){
				myShares.add((VerifiableShare) net.receiveObject());
				System.out.println("Added Share With Index " + myShares.get(i).getIndex());
			}
			
		} catch (Exception e) {
			System.out.println("Error: Could not receive share from dealer");
			System.exit(1);
		}

		System.out.println("*** SHARE VERIFICATION ***");
		int success = 1;
		for(int i = 0; i < myShares.size(); ++i){
			boolean v = myShares.get(i).verifyShare();
			System.out.println("Share " + i + " Verification: " + v);
			if(!v)
				success = -1;
			else{
				Share tmp = new Share(myShares.get(i).getShare().getS(),
						myShares.get(i).getIndex(),myShares.get(i).getShare().getParentIndex());
				tmp.setR(myShares.get(i).getR());
				tmp.setA(myShares.get(i).getA());
				tmp.setLambda(myShares.get(i).getShare().getLambda());
				
				old_store.update(myShares.get(i).getShare());
				System.out.println("Added Share With Index " + myShares.get(i).getShare().getIndex() + " to Old_Store");
				//old_store.update(tmp);
			}
		}
		if(success == 1){
			System.out.println("*** SHARE VERIFICATION COMPLETE ***");
			try{
				BufferedWriter outputStream = 
					new BufferedWriter(new FileWriter("data/" + local_hostname + "." + v++ + ".store"));
				outputStream.write(old_store.toDisk());
				outputStream.close();
			}catch(Exception e){
				System.out.println("Error writing store to disk");
			}
		}else
			System.out.println("*** SHARE VERIFICATION FAILED ***");
	}


	/**
	 * Open sockets for communication with all other servers
	 */
	public void ConnectToAll(){

		System.out.println("*** CONNECTING TO SERVERS ***");

		ArrayList<String> servers = FileIO.readFile(filename);

		// GET LOCAL ID (XINUXX... XX is ID)
		my_id = -1;
		for(int i = 0; i < n; ++i){
			if(servers.get(i).equals(local_hostname)){
				my_id = i;
				break;
			}
		}

		//Connect to Lower IDs
		for(int i = 0; i < my_id; ++i){
			try{

				Socket s = net.sendConnect(servers.get(i));
				connections.add(new ConnectionHandler(servers.get(i), i, s,
						net.getDIS(),net.getDOS(),net.getOIS(),net.getOOS()));

				System.out.println("Connected To " + servers.get(i));
			}catch(Exception e){
				e.printStackTrace();
			}
		}	

		//Wait for Connections from Higher IDs
		for(int i = my_id+1; i < n; ++i){
			try{
				Socket s = net.openConnection();
				String h = s.getInetAddress().getHostName();
				connections.add(new ConnectionHandler(h, i, s,
						net.getDIS(),net.getDOS(),net.getOIS(),net.getOOS()));
				System.out.println("Connected To " + servers.get(i));
			}catch(Exception e){
				e.printStackTrace();
			}
		}
		System.out.println("*** ALL CONNECTED ***\n");
	}

	public void VerifyStage(){

		System.out.println("*** VERIFICATION STAGE ***");


		/**
		 * Generate Subsharings for all of my Subshares
		 */
		ServerSets tmpset = new ServerSets(n,k);
		int l = tmpset.getL();
		System.out.println("L is: " + l);
		ArrayList<ArrayList<VerifyMessage>> verifyMsgs = new ArrayList<ArrayList<VerifyMessage>>();
		for(int i = 0; i < n; ++i)
			verifyMsgs.add(new ArrayList<VerifyMessage>());
		
		boolean correct = false;
		
		for (int e = 0; e < old_store.getShares().size(); ++e) { //For each share
			Share share = old_store.getShares().get(e); //This is the parent share now
			ParentIndex pi = new ParentIndex(my_id, share.getIndex(), share.getParentIndex());
			ArrayList<BigInteger> alist;
			if(share.getLambda() == null)
				alist = share.getA();
			else
				alist = share.getLambda();
			
			
			//Do sub sharing
			//We use the same random becuase we don't care about verification here
			Generator tmpG = new Generator(share.getS(),share.getR() , l);
			VerifiableSharing tmpVs = tmpG.generateVerifiableSharing();
			ArrayList<BigInteger> tmpSh = tmpVs.getS();
			ArrayList<BigInteger> tmpR = tmpVs.getR();
			ArrayList<Share> tmpShareList = new ArrayList<Share>();
			for(int j = 0; j < tmpSh.size(); j ++) {
				Share tmpshare = new Share(tmpSh.get(j),j,pi);
				tmpshare.setR(tmpR.get(j));
				tmpshare.setLambda(tmpVs.getA());
				tmpshare.setA(alist);
				tmpShareList.add(tmpshare);
			}
			
			IndexSets is = new IndexSets();
			//We use the same server sets here
			ArrayList<ArrayList<Share>> tmpServerShares = is.generateShareSets2(n,tmpShareList ,tmpset.constructServerSets());
			//Store the server shares in each sharestore
			for(int j = 0; j < n; j++) {
				//jth server
				ArrayList<Share> tmpShList = tmpServerShares.get(j); //The share list for server j
				for (Iterator<Share> it = tmpShList.iterator(); it.hasNext();) {
					
					//Making a copy of the share
					Share primarytmp = it.next();
					Share tmp = new Share(primarytmp.getS(),primarytmp.getIndex(),primarytmp.getParentIndex());
					tmp.setA(primarytmp.getA());
					tmp.setR(primarytmp.getR());
					tmp.setLambda(primarytmp.getLambda());
					
					int host = -1;
					if(j == my_id){
						new_store.update(tmp);
						continue;
					}
					for(int u = 0; u < connections.size(); ++u)
						if(connections.get(u).getIndex() == j){
							host = u;
							break;
						}
					
					
					crypto.genNewAESKey();
					
					VerifiableShare _tmp 
					= new VerifiableShare(
							tmp,
							tmp.getA(),
							tmp.getA().get(0),tmp.getIndex(),
							pi,
							crypto.encrypt(crypto.getAESKey(),connections.get(host).getHost()));
					
					_tmp.setEncS(crypto.aes_encrypt(
							tmp.getS().toByteArray()));

					_tmp.setEncR(crypto.aes_encrypt(
							tmp.getR().toByteArray()));
					
					tmp.setR(null);
					tmp.setS(null);

					VerifyMessage v = new VerifyMessage(_tmp);

					v.destination = connections.get(host).getHost();
					v.source = local_hostname;
					
					byte[] signature = BigInteger.ONE.toByteArray();
					if(!failures){
						signature = crypto.sign(objectToBytes(v));
					}
					
					v.setSignature(signature);

					verifyMsgs.get(j).add(v);
					
					
				}
			}
		}
		
		System.out.println("2) Verify Messages Constructed");
		
		l = 0;
		for(int i = 0; i < verifyMsgs.size(); ++i)
			if(verifyMsgs.get(i).size() > l)
				l = verifyMsgs.get(i).size();
		
		System.out.println("2.1) Sending " + l + " Verify Messages to each server");


		/**
		 * Send Verify Messages
		 */
		if(!nonresponsive){
			sendVerifyMsgs(verifyMsgs);
			System.out.println("3) Verify Messages Sent");
		}else{
			System.out.println("6) Nonresponsive: No Verify Messages Sent");
		}


		/**
		 * Receive Verify Messages
		 */
		ArrayList<VerifyMessage> recVerifyMsgs = new ArrayList<VerifyMessage>();
		receiveVerifyMsgs(recVerifyMsgs,l);
		System.out.println("4) Received " + recVerifyMsgs.size() + " Verify Messages");
		System.out.println("4.1) Verified Servers Size: " + verifiedServers.size());
		System.out.println("T+1: " + (t+1));
		for(int i = 0; i < verifiedServers.size(); ++i)
			System.out.println("Verified Server: " + verifiedServers.get(i));
		if(verifiedServers.size() < (t+1)){
			System.out.println("*** VERIFICATION STAGE INCOMPLETE ***");
			System.out.println("*** Less than T+1 Verified Messages Received ***");
		}
		
		/**
		 * Verify Messages and send Verified Messages 
		 */
		ArrayList<VerifiedMessage> verified = new ArrayList<VerifiedMessage>();
		buildVerifiedMessages(recVerifyMsgs,verified);
		System.out.println("5) Verify Message Verification Complete");

		/**
		 * Send Verified Messages
		 */
		if(!nonresponsive){
			sendVerifiedMsgs(verified);
			System.out.println("6) Verified Messages Sent");
		}else{
			System.out.println("6) Nonresponsive: No Verified Messages Sent");
		}

		/**
		 * Receive Verified Messages
		 */
		ArrayList<VerifiedMessage> recVerified = new ArrayList<VerifiedMessage>();
		int numVerified = receiveVerifiedMsgs(recVerified,l);
		System.out.println("7) Received " + numVerified + " Verified Messages");

		if(verifiedServers.size() >= t+1){
			/*
			 * We have enough verified messages to update the ShareStores
			 */
			try{
				BufferedWriter outputStream = 
					new BufferedWriter(new FileWriter("data/" + local_hostname + "." + v++ + ".store"));
				outputStream.write(new_store.toDisk());
				outputStream.close();
			}catch(Exception e){
				System.out.println("Error writing store to disk");
			}
			
			System.out.println("Store Size: " + new_store.getShares().size());
			
			old_store = new_store;
			new_store = null;
			new_store = new ShareStore();
			

			System.out.println("*** VERIFICATION STAGE COMPLETE ***");
		}else{
			System.out.println("*** VERIFICATION STAGE INCOMPLETE ***");
			System.out.println("*** Less than T+1 Verified Messages Received ***");
			System.exit(1);
		}
	}

	public BigInteger RecoverStage(){

		BigInteger reconstructed = null;


		return reconstructed;
	}


	public int receiveVerifiedMsgs(ArrayList<VerifiedMessage> recVerified, int l){
		
		for(int i = 0; i < verifiedServers.size(); ++i) {
			
			int index = -1;
			
			if(verifiedServers.get(i).equals(local_hostname)){
				continue;
			}
			
			try{
				
				index = 0;
				for(int j = 0; j < connections.size(); ++j){
					if(verifiedServers.get(i).equals(connections.get(j).getHost())){
						index = j;
						break;
					}
				}
				
				net.changeConnection(connections.get(index));
				net.setSocketTimeout(0);
			}catch(Exception e){
				System.out.println("Error Setting up Network Communicator");
			}

			for(int j = 0; j < l; ++j){
				try{

					VerifiedMessage v = (VerifiedMessage) net.receiveObject();

					if(verifyDigitalSignatureVerified(v)){
						recVerified.add(v);
					}else{
						System.out.println("Verified Message Digital Signature Failed");
						j = l;
					}

				}catch(SocketTimeoutException t){
					System.out.println("Connection to host " + connections.get(index).getHost() + " timed out");
					j = l;
				}catch(Exception e){
					System.out.println("Error reading Verified Message from " + connections.get(index).getHost());
					j = l;
				}
			}
		}

		return recVerified.size();
	}

	/**
	 * Constructs the Verified messages for all Verify messages
	 * that are successfully verified
	 * @param verifyMsgs
	 * @param verifiedMsgs
	 */
	public void buildVerifiedMessages(ArrayList<VerifyMessage> verifyMsgs, 
			ArrayList<VerifiedMessage> verifiedMsgs){

		for(int i = 0; i < verifyMsgs.size(); ++i){

			VerifyMessage v = verifyMsgs.get(i);

			if(verifyMessage(v)){

				new_store.update(v.getShare().getShare());

				VerifiedMessage tmp = new VerifiedMessage(v.getA(),v.getLambda(),v.getIndex());
				
				tmp.source = local_hostname;
				tmp.destination = v.source;

				tmp.setSignature(crypto.sign(objectToBytes(tmp)));

				verifiedMsgs.add(tmp);
			}
		}
	}


	/**
	 * Send Verified Messages
	 * @param verifiedMsgs
	 */
	public void sendVerifiedMsgs(ArrayList<VerifiedMessage> verifiedMsgs){

		ArrayList<String> badHosts = new ArrayList<String>();
		for(int i = 0; i < verifiedMsgs.size(); ++i){
			int index = -1;
			try {

				
				for(int k = 0; k < connections.size(); ++k){
					if(connections.get(k).getHost().equals(verifiedMsgs.get(i).destination)){
						index = k;
						break;
					}
				}
				
				if(badHosts.contains(connections.get(index).getHost())){
						continue;
				}

				net.changeConnection(connections.get(index));
				net.setSocketTimeout(0);
				net.sendObject(verifiedMsgs.get(i));
			} catch(SocketTimeoutException t){
				System.out.println("Connection to host " + connections.get(index).getHost() + " timed out");
			}catch (Exception e) {
				System.out.println("Error Sending to Host " + connections.get(index).getHost());
				badHosts.add(connections.get(index).getHost());
			}
		}
	}


	/**
	 * Converts an object into a byte array
	 * @param o
	 * @return
	 */
	public byte[] objectToBytes(Object o){

		ByteArrayOutputStream bos = new ByteArrayOutputStream();

		try{
			ObjectOutput out = new ObjectOutputStream(bos);
			out.writeObject(o);
		}catch(Exception e){
			e.printStackTrace();
			System.exit(1);
		}

		return bos.toByteArray();
	}


	/**
	 * Send an array of verify messages
	 * @param verifyMsgs
	 */
	public void sendVerifyMsgs(ArrayList<ArrayList<VerifyMessage>> verifyMsgs){

		if(failures){
			
			double rand = Math.random();
			if(rand < 0.5){
				System.out.println("Malicious: Dropping Message");
				return;
			}
		}
		
		int index = -1;
		for(int i = 0; i < verifyMsgs.size(); ++i){
			for(int j = 0; j < verifyMsgs.get(i).size(); ++j){
				try {

					
					for(int k = 0; k < connections.size(); ++k){
						if(connections.get(k).getHost().equals(verifyMsgs.get(i).get(j).destination)){
							index = k;
							break;
						}
					}

					net.changeConnection(connections.get(index));
					net.setSocketTimeout(0);
					net.sendObject(verifyMsgs.get(i).get(j));
				} catch(SocketTimeoutException t){
					System.out.println("Connection to host " + connections.get(index).getHost() + " timed out");
					j = verifyMsgs.get(i).size();
				}catch (Exception e) {
					System.out.println("Error Sending to Host " + connections.get(index).getHost());
					j = verifyMsgs.get(i).size();
				}
			}
		}
	}


	/**
	 * Receive Verify Messages from other servers
	 * @param rec
	 * @param l
	 */
	public void receiveVerifyMsgs(ArrayList<VerifyMessage> rec, int l){

		ArrayList<String> badServers = new ArrayList<String>();
		ArrayList<VerifyMessage> ret = new ArrayList<VerifyMessage>();
		
		for(int i = 0; i < n-1; ++i) {
			try{
				net.changeConnection(connections.get(i));
			}catch(Exception e){
				System.out.println("Error Setting up Network Communicator");
			}

			for(int j = 0; j < l; ++j){
				try{

					net.setSocketTimeout(timeout);
					VerifyMessage v = (VerifyMessage) net.receiveObject();
					//System.out.println("Received VM from " + v.source);
					if(verifyDigitalSignature(v)){
						//System.out.println("Valid DS VM from " + v.source);
						decryptVerifyMsg(v);
						if(verifyMessage(v)){
							//System.out.println("Adding VM from " + v.source);
							ret.add(v);
							if(!verifiedServers.contains(v.source))
								verifiedServers.add(v.source);
						}
					}else{
						System.out.println("Digital Signature Failed!");
						if(!badServers.contains(v.source))
							badServers.add(v.source);
					}

				}catch(SocketTimeoutException t){
					System.out.println("Connection to host " + connections.get(i).getHost() + " timed out");
					j = l;
				}catch(Exception e){
					System.out.println("Error Reading Verify Message from Host " + connections.get(i).getHost());
					j = l;
				}
			}
		}
		
		
		for(int i = 0; i < ret.size(); ++i){
			String host = ret.get(i).source;
			if(!badServers.contains(host)){
				rec.add(ret.get(i));
			}else{
				System.out.println("Removing Message from " + ret.get(i).source);
			}
		}
	}


	/**
	 * Extracts the AES Key and decrypts S and R
	 * @param v
	 * @return
	 */
	public VerifyMessage decryptVerifyMsg(VerifyMessage v){

		byte[] decAesKey = crypto.decrypt(v.getShare().getAesRaw());
		crypto.setAESKey(decAesKey);

		v.getShare().setS(
				new BigInteger(crypto.aes_decrypt(v.getShare().getEncS())));
		v.getShare().setR(
				new BigInteger(crypto.aes_decrypt(v.getShare().getEncR())));

		return v;
	}

	/**
	 * Returns TRUE if V has a valid Digital Signature by Source
	 * @param v
	 * @return
	 */
	public boolean verifyDigitalSignature(VerifyMessage v){
		byte[] sig = v.getSignature();
		v.setSignature(null);

		ByteArrayOutputStream bos = new ByteArrayOutputStream();

		try{
			ObjectOutput out = new ObjectOutputStream(bos);
			out.writeObject(v);
		}catch(Exception e){
			e.printStackTrace();
			System.exit(1);
		}

		boolean valid = crypto.verifySignature(sig, bos.toByteArray(), 
				v.source);

		v.setValidSig(valid);

		v.setSignature(sig);

		return valid;
	}


	/**
	 * Returns TRUE if V has a valid Digital Signature by Source
	 * @param v
	 * @return
	 */
	public boolean verifyDigitalSignatureVerified(VerifiedMessage v){
		byte[] sig = v.getSignature();
		v.setSignature(null);

		ByteArrayOutputStream bos = new ByteArrayOutputStream();

		try{
			ObjectOutput out = new ObjectOutputStream(bos);
			out.writeObject(v);
		}catch(Exception e){
			e.printStackTrace();
			System.exit(1);
		}

		boolean valid = crypto.verifySignature(sig, bos.toByteArray(), 
				v.source);

		v.setSigValid(valid);

		v.setSignature(sig);

		return valid;
	}

	/**
	 * Returns TRUE if V satisfies all necessary conditions for validation
	 * @param v
	 * @return
	 */
	public boolean verifyMessage(VerifyMessage v){

		Verifier ver = new Verifier();
		v.setSignature(null);

		ByteArrayOutputStream bos = new ByteArrayOutputStream();

		try{
			ObjectOutput out = new ObjectOutputStream(bos);
			out.writeObject(v);
		}catch(Exception e){
			e.printStackTrace();
			System.exit(1);
		}
		
		Share s = v.getShare().getShare();
		
		for(int i = 1; i < s.getA().size(); ++i)
			if(s.getA().get(i).equals(s.getLambda().get(0)))
				if(i != s.getParentIndex().getIndex()+1)
					System.out.println("Need Index " + i + " Rather Than " + s.getParentIndex().getIndex()+1);
		

		//System.out.println("Verify Message From " + v.source);

		boolean digSig = v.getSigValid();
		if(!digSig)
			System.out.println("Verify Digital Signature: " + digSig);
		boolean verifyA = ver.verifyVecStandAlone(s.getA());
		if(!verifyA)
			System.out.println("Verify Vec A: " + verifyA);
		boolean verifyL =  ver.verifyVecStandAlone(s.getLambda());
		if(!verifyL)
			System.out.println("Verify Vec L: " + verifyL);
		boolean verifyAL = s.getA().get(s.getParentIndex().getIndex()+1).equals(s.getLambda().get(0));
		if(!verifyAL)
			System.out.println("Verify A[i] == Li[0]: " + verifyAL);
		boolean verifyShare = ver.verifyShareClass(v.getShare());
		if(!verifyShare)
			System.out.println("Verify Share: " + verifyShare);

		if(digSig && verifyA && verifyL && verifyAL && verifyShare)
			return true;
		return false;
	}

}
